# -*- coding: utf-8 -*-
# @Time    : 2020/3/30 21:07
# @Author  : zhangyedi
# @Email: 178320369@qq.com
# @File:machine_learn.py
""" 构建决策树房价预测模型 """

from sklearn.tree import DecisionTreeRegressor
from sklearn.neighbors import KNeighborsRegressor
import pandas as pd

# 1. 设置行不限制数量
pd.set_option('display.max_rows', None)
# 2. 设置列不限制数量
pd.set_option('display.max_columns', None)

df = pd.read_csv("ml.csv")
df = df.dropna(axis=0)  # 删除错误样本
# print(df.head())
# 数据处理
df["direction"].replace(["东", "南", "西", "北"], [1, 2, 3, 4], inplace=True)  # 用1234表示东南西北
# print(df["renovation"].unique()) #['精装 ' '简装 ' '毛坯 ']
df["renovation"].replace(['精装 ', '简装 ', '毛坯 '], [1, 2, 3], inplace=True)  # 用1,2,3['精装 ' '简装 ' '毛坯 ']
# print(df["ta"].unique()) #['塔楼' '板塔结合' '板楼' '平房' '暂无数据']
df["ta"].replace(['塔楼', '板塔结合', '板楼', '平房', '暂无数据'], [1, 2, 3, 4, 5],
                 inplace=True)  # 用1,2,3,4,5['塔楼' '板塔结合' '板楼' '平房' '暂无数据']
# print(df.head())
# 1、构造数据
y = df["totalPrice"].astype(float)  # 房价;标签集
# print(y)
x = df.drop(["totalPrice", "unitPrice"], axis=1)

# print(x)
# 2、构造模型
# model = KNeighborsRegressor()
model = DecisionTreeRegressor()
# 3、模型训练
model.fit(x, y)
# 4、模型预测
# 黄同学说：90平方米，朝南，精装修，第8层，平房，2室2厅，位置佛山
huang_house = [[100, 2, 1, 1, 4, 6, 6]]
totalprice = model.predict(huang_house)

print(f"模型预测张业頔同学心仪的房子总价为：{totalprice}万元")
